//package com.fy.fyai.common.config;
//
//import com.fy.fyai.service.AiChatConversationService;
//import org.reactivestreams.Publisher;
//import org.springframework.ai.chat.client.advisor.AroundAdvisor;
//import org.springframework.ai.chat.client.advisor.ChatClientRequestContext;
//import org.springframework.ai.chat.metadata.Usage;
//import org.springframework.ai.chat.model.ChatResponse;
//import org.springframework.stereotype.Component;
//import reactor.core.publisher.Flux;
//import reactor.core.publisher.Mono;
//
//import java.util.Map;
//import java.util.concurrent.atomic.AtomicInteger;
//
//@Component
//public class TokenUsageAdvisor implements AroundAdvisor {
//
//    private final AiChatConversationService aiChatConversationService;
//
//    public TokenUsageAdvisor(AiChatConversationService aiChatConversationService) {
//        this.aiChatConversationService = aiChatConversationService;
//    }
//
//    @Override
//    public Mono<ChatResponse> aroundCall(Mono<ChatResponse> mono, ChatClientRequestContext ctx) {
//        Map<String, Object> params = ctx.getAdvisorParams();
//        String conversationId = String.valueOf(params.getOrDefault("chat_memory_conversation_id", "default"));
//
//        return mono.doOnNext(resp -> {
//            Usage u = (resp.getMetadata() != null) ? resp.getMetadata().getUsage() : null;
//            if (u != null) {
//                aiChatConversationService.saveTokens(
//                        conversationId,
//                        u.getPromptTokens(),
//                        u.getCompletionTokens(),
//                        u.getTotalTokens()
//                );
//            }
//        });
//    }
//
//    @Override
//    public Publisher<ChatResponse> aroundStreamCall(Publisher<ChatResponse> publisher, ChatClientRequestContext ctx) {
//        Map<String, Object> params = ctx.getAdvisorParams();
//        String conversationId = String.valueOf(params.getOrDefault("chat_memory_conversation_id", "default"));
//
//        AtomicInteger prompt = new AtomicInteger();
//        AtomicInteger completion = new AtomicInteger();
//        AtomicInteger total = new AtomicInteger();
//
//        return Flux.from(publisher)
//                .doOnNext(resp -> {
//                    Usage u = (resp.getMetadata() != null) ? resp.getMetadata().getUsage() : null;
//                    if (u != null) {
//                        prompt.addAndGet(u.getPromptTokens());
//                        completion.addAndGet(u.getCompletionTokens());
//                        total.addAndGet(u.getTotalTokens());
//                    }
//                })
//                .doOnComplete(() -> aiChatConversationService.saveTokens(
//                        conversationId, prompt.get(), completion.get(), total.get()
//                ));
//    }
//}